

#libraries and modules
import os
from array import *
import csv
import math
import random
import numpy as np
import numpy.ma as ma
import scipy as sp
import scipy.stats.stats as st
from scipy.interpolate import UnivariateSpline
from scipy.optimize import curve_fit
from scipy.ndimage.filters import generic_filter
import scipy.stats.stats as st
from scipy.stats import t,norm,linregress,nanmean#,nansum
from lmfit import minimize, Parameters,report_errors

import matplotlib.pyplot as plt
import matplotlib.mlab as mlab

from matplotlib.ticker import MaxNLocator, AutoMinorLocator
from matplotlib import transforms as mtransforms
from matplotlib import font_manager as font
from matplotlib.figure import Figure #as mFig
from matplotlib.axis import Axis



#math/trigo functions
tan=math.tan
sin=math.sin
asin=math.asin
cos=math.cos
pow=math.pow
sqrt=math.sqrt
atan=math.atan
deg=math.degrees
rad=math.radians
log=math.log10
exp=math.exp


def generate_CN_dist(minCN,maxCN):
    CN=np.random.random(1000)*(maxCN-minCN)+minCN
    return CN

def compute_S(CN):
    S20=(1000/(1.*CN))-10
    #S05=1.33*S20**1.15   #from Hawkins, R.H.; Jiang, R.; Woodward, D.E.; Hjelmfelt, A.T.; Van Mullem, J.A. (2002). "Runoff Curve Number Method: Examination of the Initial Abstraction Ratio". Proceedings of the Second Federal Interagency Hydrologic Modeling Conference, Las Vegas, Nevada (U.S. Geological Survey). Retrieved 24 November 2013.
    return S20

def compute_Q(P,S):
    #Q=(P-Ia)**2 / ((P-Ia)+S); assuming Ia=0.05S;
    P=P/25.4
    Q=np.where(P<=0.2*S,
               np.zeros(S.shape),
               25.4*((P-0.2*S)**2)/(P+0.8*S))
    return Q

def compute_Tsheet(n,L,P2,s):
    L=L*100/(2.54*12)
    # n Manning's coefficient
    # L flow length (ft)
    # P2 2-year 24-hr rainfall (in)
    # s=land slope (ft/ft)
    Tsheet=60*(0.007*(n*L)**0.8)/((P2**0.5)*s**0.4) #in minutes
    smd,mean,std=SMD_analysis(Tsheet,100,30,0.05)
    return smd,mean,std
    


def samplemeandist(x,numsamples,samp_size):
    sample_means=array('f')
    for a in range(numsamples):
        indices = np.arange(len(x))
        np.random.shuffle(indices)
        samp_indices=indices[:samp_size]
        sample=x[samp_indices]
        sample_means.append(np.mean(sample))
    return sample_means



def KS2samp_normtest(x,alpha):
    normal_samp=np.random.normal(np.mean(x),np.std(x),len(x)*3)
    KS,pKS=st.ks_2samp(x, normal_samp)
    if pKS>alpha:return "normal"
    else:return "not normal"



def SMD_analysis(x,numsamples,samp_size,alpha):
    #creating sample mean distribution
    cur_smd=samplemeandist(x,numsamples,samp_size)

    #checking if sample means are normally distributed
    #print KS2samp_normtest(cur_smd,alpha)
    
    #6.3 plotting fitted normal distribution curve
    mu, std = norm.fit(cur_smd)
    x = np.linspace(min(cur_smd),max(cur_smd), 200)
    p = norm.pdf(x, mu, std)
    #plt.plot(x, p, '-',linewidth=2, label=cur_scen_label)
    
        
    
    
    return cur_smd,mu,std

plt.ion()
plt.close("all")

#1. Loss calculation (using SCS-CN)
Pre_CN=(63,80)
Post_CN=(80,93)
scenarioCN=(Pre_CN,Post_CN)

#2. Concentration time calculation
#Tsheet=(0.007*(n*L)**0.8)/(P2**0.5)*s*0.4
minP2=50/25.4  #max 2-year 24-hr rainfall, inches
maxP2=300/25.4 #max 2-year 24-hr rainfall, inches
L=270    #flowlength, (m)
         #n, manning roughness coefficient, table 3-1, TR55
Pre_TC=(0.1,0.15)    #(s=natural,n=shortgrass prairie)
Post_TC=(0.05,0.02)  #(s=graded,n=conservatively high value for smooth surfaces)
scenarioTC=(Pre_TC,Post_TC)




scenario_label=('GRASSLAND','BUILT-UP')

CN_list=[]
TC_list=[]
tlag_list=[]
for a in range(2):
    #LOSS 
    cur_scen=scenarioCN[a]
    cur_scen_label=scenario_label[a]
    print "\n"+cur_scen_label
    #plt.figure(1)
    CN=generate_CN_dist(cur_scen[0],cur_scen[1])
    smd,mean,std=SMD_analysis(CN,100,30,0.05)
    
    if a==0:
        print "CN (95%): ",round(mean+2*std,1)
        CN_list.append(round(mean+2*std,1))

    else:
        print "CN (05%): ",round(mean-2*std,1)
        CN_list.append(round(mean-2*std,1))
        #plt.legend()
        #plt.xlabel("SCS Curve number")
    #CONCENTRATION TIME
    P2=np.random.random(1000)*(maxP2-minP2)+minP2
    cur_scen=scenarioTC[a]
    cur_s=cur_scen[0]
    cur_n=cur_scen[1]
    #plt.figure(2)
    smd,mean,std=compute_Tsheet(cur_n,L,P2,cur_s)
    if a==0:
        print "Tc: ",round(mean-2*std,1)
        TC_list.append(round(mean-2*std,1))
    else:
        print "Tc: ",round(mean+2*std,1)
        TC_list.append(round(mean+2*std,1))
        #plt.xlabel("Concentration time, min")
        #plt.legend

    
#plt.show()  







# 3. hydrograph analysis
#basin characteristics
A=270*100  #area of basin, m2
S_list=compute_S(np.asarray(CN_list))



#rainfall model
Plist=[2.5,7.5,15,30]#,100,150,200]

beforepeakQ=[]
afterpeakQ=[]
beforelag=[]
afterlag=[]
for Ptotal in Plist:
    #Ptotal=150  #mm
    Pduration=60    #min
    Ppeak=Ptotal*2/(Pduration*1.)  #mm, peak rainfall, assuming triangular rainfall distribution
    Ppeakloc=Pduration/3.  #  #time to peak q
        

    #fig,axes=plt.subplots(nrows=3,ncols=1,sharex=True)
            

    peakQ=[]
    peakTlag=[]
    print "\n\n"
    print "Total rainfall depth, mm: ",Ptotal
    print "Total rainfall volume, m3: ",A*Ptotal/1000.
    print "Rainfall duration, min: ", Pduration
    print "Peak rainfall, mm: ", Ppeak
        
    for t in range(2):
        fig,axes=plt.subplots(nrows=3,ncols=1,sharex=True)
        plt.suptitle(scenario_label[t]+", "+str(Ptotal)+"mm")
        print "\n"+scenario_label[t]
        tc=TC_list[t]  #minutes
        CN=CN_list[t]
        S=S_list[t]
        print "CN: ",CN,"\nS: ",round(S,1),"\nTc: ",tc
            
        #UH parameters
        tlag=0.6*tc
        dtime=0.2*0.6*min(TC_list)
        Tp=0.5*dtime+tlag
        Tb=2.67*Tp
        Up=2.08*A/Tp

        #print round(dtime,2),round(Tp,2)

        #rainfall model
        m1=(Ppeak-0)/(Ppeakloc-0)
        b1=Ppeak-m1*Ppeakloc
        m2=(0-Ppeak)/(Pduration-Ppeakloc)
        b2=Ppeak-m2*Ppeakloc
        
        model_time=dtime*np.arange(int(5*Pduration/dtime))#np.linspace(0,10*Pduration,int(10*Pduration/dtime))
        model_p=np.where(model_time<=Ppeakloc,
                         m1*model_time+b1,
                         m2*model_time+b2)
        model_p=np.where(model_p>0,model_p,np.zeros(model_p.shape))
        model_ex_p=compute_Q(model_p,S)
        
        axes[0].bar(model_time,model_p,width=dtime,facecolor='b')
        axes[0].set_ylabel('rainfall,\nmm')
        axes[1].bar(model_time,model_ex_p,width=dtime,facecolor='r')
        axes[1].set_ylabel('excess rainfall,\nmm')


        #plt.show()

        q=np.zeros(model_time.shape)
        for h in range(len(model_ex_p)):
            cur_ex_p=model_ex_p[h]
            if cur_ex_p<=0:continue
            cur_time=model_time[h]

            curT0=cur_time
            curTp=Tp+curT0
            curTb=Tb+curT0

            cur_base=curTb-curT0
            cur_Up=(cur_ex_p/1000.)*A*2/cur_base
            #print h,cur_time,cur_ex_p,cur_Up,curTb-curT0
          
            m1=(cur_Up-0)/(curTp-curT0)
            b1=cur_Up-m1*curTp

            m2=(0-cur_Up)/(curTb-curTp)
            b2=cur_Up-m2*curTp

            cur_q=np.where(model_time<curT0,
                           np.zeros(model_time.shape),
                           np.where(model_time<curTp,
                                    m1*model_time+b1,
                                    np.where(model_time<curTb,
                                             m2*model_time+b2,
                                             np.zeros(model_time.shape))))
                                    

            #axes[2].plot(model_time,cur_q/60.,'-')

            
            for i in range(h,len(model_ex_p)):
                if model_time[i]<=curTp:
                    q[i]=q[i]+(m1*model_time[i]+b1)
                else:
                    if m2*model_time[i]+b2>=0:
                        q[i]=q[i]+(m2*model_time[i]+b2)
                    else:
                        q[i]=q[i]+0
                        continue
                        
        q=np.where(q>0,q,np.zeros(model_time.shape))
        axes[2].bar(model_time,q/60.,width=dtime,facecolor='g')
        axes[2].set_ylabel('discharge,m3/s')
        axes[2].set_xlabel('time, min')
        print "peak basin discharge: ",round(max(q),3), "m3/s"
        if round((model_time[np.argmax(q)]-model_time[np.argmax(model_p)]),2)<0:
            print "rainfall-discharge peak lag: 0.0"
        else:
            print "rainfall-discharge peak lag: ",round((model_time[np.argmax(q)]-model_time[np.argmax(model_p)]),2), "min"
        
        if t==0:
            beforepeakQ.append(round(max(q),3))
            if round((model_time[np.argmax(q)]-model_time[np.argmax(model_p)]),2)<0:
                beforelag.append(np.nan)
            else:
                beforelag.append(round((model_time[np.argmax(q)]-model_time[np.argmax(model_p)]),2))
        else:
            afterpeakQ.append(round(max(q),3))
            if round((model_time[np.argmax(q)]-model_time[np.argmax(model_p)]),2)<0:
                afterlag.append(np.nan)
            else:
                afterlag.append(round((model_time[np.argmax(q)]-model_time[np.argmax(model_p)]),2))
    
    #plt.show()
plt.figure()
plt.plot(Plist,beforepeakQ,'o-',label='grassland')
plt.plot(Plist,afterpeakQ,'o-',label='built-up')
plt.title("Peak basin discharge for 60-min rainfall")
plt.ylabel("discharge, m3/s")
plt.xlabel("Total rainfall depth, mm")
plt.legend()

plt.figure()
plt.plot(Plist,beforelag,'o-',label='grassland')
plt.plot(Plist,afterlag,'o-',label='built-up')
plt.title("rainfall-discharge peak lag for 60-min rainfall")
plt.ylabel("lag time, min")
plt.xlabel("Total rainfall depth, mm")
plt.legend()




plt.show()



    


